//! Predicates for disabling compression of responses.
//!
//! Predicates are applied with [`Compression::compress_when`] or
//! [`CompressionLayer::compress_when`].
//!
//! [`Compression::compress_when`]: super::Compression::compress_when
//! [`CompressionLayer::compress_when`]: super::CompressionLayer::compress_when

use rama_core::{extensions::Extensions, extensions::ExtensionsRef};
use rama_http_types::{HeaderMap, StatusCode, StreamingBody, Version, header};
use rama_utils::str::arcstr::{ArcStr, arcstr};

/// Predicate used to determine if a response should be compressed or not.
pub trait Predicate: Clone {
    /// Should this response be compressed or not?
    fn should_compress<B>(&self, response: &rama_http_types::Response<B>) -> bool
    where
        B: StreamingBody;

    /// Combine two predicates into one.
    ///
    /// The resulting predicate enables compression if both inner predicates do.
    fn and<Other>(self, other: Other) -> And<Self, Other>
    where
        Self: Sized,
        Other: Predicate,
    {
        And {
            lhs: self,
            rhs: other,
        }
    }
}

impl<F> Predicate for F
where
    F: Fn(StatusCode, Version, &HeaderMap, &Extensions) -> bool + Clone,
{
    fn should_compress<B>(&self, response: &rama_http_types::Response<B>) -> bool
    where
        B: StreamingBody,
    {
        let status = response.status();
        let version = response.version();
        let headers = response.headers();
        let extensions = response.extensions();
        self(status, version, headers, extensions)
    }
}

impl<T> Predicate for Option<T>
where
    T: Predicate,
{
    fn should_compress<B>(&self, response: &rama_http_types::Response<B>) -> bool
    where
        B: StreamingBody,
    {
        self.as_ref()
            .map(|inner| inner.should_compress(response))
            .unwrap_or(true)
    }
}

/// Two predicates combined into one.
///
/// Created with [`Predicate::and`]
#[derive(Debug, Clone, Default, Copy)]
pub struct And<Lhs, Rhs> {
    lhs: Lhs,
    rhs: Rhs,
}

impl<Lhs, Rhs> Predicate for And<Lhs, Rhs>
where
    Lhs: Predicate,
    Rhs: Predicate,
{
    fn should_compress<B>(&self, response: &rama_http_types::Response<B>) -> bool
    where
        B: StreamingBody,
    {
        self.lhs.should_compress(response) && self.rhs.should_compress(response)
    }
}

/// The default predicate used by [`Compression`] and [`CompressionLayer`].
///
/// This will compress responses unless:
///
/// - They're gRPC, which has its own protocol specific compression scheme.
/// - It's an image as determined by the `content-type` starting with `image/`.
/// - They're Server-Sent Events (SSE) as determined by the `content-type` being `text/event-stream`.
/// - The response is less than 32 bytes.
///
/// # Configuring the defaults
///
/// `DefaultPredicate` doesn't support any configuration. Instead you can build your own predicate
/// by combining types in this module:
///
/// ```rust
/// use rama_utils::str::arcstr::arcstr;
/// use rama_http::layer::compression::predicate::{SizeAbove, NotForContentType, Predicate};
///
/// // slightly large min size than the default 32
/// let predicate = SizeAbove::new(256)
///     // still don't compress gRPC
///     .and(NotForContentType::GRPC)
///     // still don't compress images
///     .and(NotForContentType::IMAGES)
///     // also don't compress JSON
///     .and(NotForContentType::new(arcstr!("application/json")));
/// ```
///
/// [`Compression`]: super::Compression
/// [`CompressionLayer`]: super::CompressionLayer
#[derive(Debug, Clone)]
pub struct DefaultPredicate(
    And<And<And<SizeAbove, NotForContentType>, NotForContentType>, NotForContentType>,
);

impl DefaultPredicate {
    /// Create a new `DefaultPredicate`.
    #[must_use]
    pub fn new() -> Self {
        let inner = SizeAbove::new(SizeAbove::DEFAULT_MIN_SIZE)
            .and(NotForContentType::GRPC)
            .and(NotForContentType::IMAGES)
            .and(NotForContentType::SSE);
        Self(inner)
    }
}

impl Default for DefaultPredicate {
    fn default() -> Self {
        Self::new()
    }
}

impl Predicate for DefaultPredicate {
    fn should_compress<B>(&self, response: &rama_http_types::Response<B>) -> bool
    where
        B: StreamingBody,
    {
        self.0.should_compress(response)
    }
}

/// [`Predicate`] that will only allow compression of responses above a certain size.
#[derive(Clone, Copy, Debug)]
pub struct SizeAbove(u16);

impl SizeAbove {
    pub(crate) const DEFAULT_MIN_SIZE: u16 = 32;

    /// Create a new `SizeAbove` predicate that will only compress responses larger than
    /// `min_size_bytes`.
    ///
    /// The response will be compressed if the exact size cannot be determined through either the
    /// `content-length` header or [`Body::size_hint`].
    #[must_use]
    pub const fn new(min_size_bytes: u16) -> Self {
        Self(min_size_bytes)
    }
}

impl Default for SizeAbove {
    fn default() -> Self {
        Self(Self::DEFAULT_MIN_SIZE)
    }
}

impl Predicate for SizeAbove {
    fn should_compress<B>(&self, response: &rama_http_types::Response<B>) -> bool
    where
        B: StreamingBody,
    {
        let content_size = response.body().size_hint().exact().or_else(|| {
            response
                .headers()
                .get(header::CONTENT_LENGTH)
                .and_then(|h| h.to_str().ok())
                .and_then(|val| val.parse().ok())
        });

        match content_size {
            Some(size) => size >= (self.0 as u64),
            _ => true,
        }
    }
}

/// Predicate that wont allow responses with a specific `content-type` to be compressed.
#[derive(Clone, Debug)]
pub struct NotForContentType {
    content_type: ArcStr,
    exception: Option<ArcStr>,
}

impl NotForContentType {
    /// Predicate that wont compress gRPC responses.
    pub const GRPC: Self = Self::new(arcstr!("application/grpc"));

    /// Predicate that wont compress images.
    pub const IMAGES: Self = Self {
        content_type: arcstr!("image/"),
        exception: Some(arcstr!("image/svg+xml")),
    };

    /// Predicate that wont compress Server-Sent Events (SSE) responses.
    pub const SSE: Self = Self::new(arcstr!("text/event-stream"));

    /// Create a new `NotForContentType`.
    #[must_use]
    pub const fn new(content_type: ArcStr) -> Self {
        Self {
            content_type,
            exception: None,
        }
    }
}

impl Predicate for NotForContentType {
    fn should_compress<B>(&self, response: &rama_http_types::Response<B>) -> bool
    where
        B: StreamingBody,
    {
        if let Some(except) = &self.exception
            && content_type(response) == except.as_str()
        {
            return true;
        }

        !content_type(response).starts_with(self.content_type.as_str())
    }
}

fn content_type<B>(response: &rama_http_types::Response<B>) -> &str {
    response
        .headers()
        .get(header::CONTENT_TYPE)
        .and_then(|h| h.to_str().ok())
        .unwrap_or_default()
}
